import asyncio
import threading
import websockets
import json
import logging
import secrets
from datetime import datetime

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class Server:
    def __init__(self):
        self.ws_thread = None
        self.running = False
        self.ws_server = None
        self.clients = {}  # {token: websocket}
        self.app_status = {}  # {token: last_heartbeat_time}

    def generate_token(self):
        """生成安全的token"""
        return secrets.token_urlsafe(32)

    async def register(self, websocket, token):
        """注册新的客户端连接"""
        self.clients[token] = websocket
        self.app_status[token] = datetime.now()
        logger.info(f"新客户端注册: {token}")

    async def unregister(self, token):
        """注销客户端连接"""
        if token in self.clients:
            del self.clients[token]
            del self.app_status[token]
            logger.info(f"客户端注销: {token}")

    def is_app_running(self, token):
        """检查应用是否在运行"""
        if token not in self.app_status:
            return False
        # 检查最后心跳时间是否在60秒内
        return (datetime.now() - self.app_status[token]).total_seconds() < 60

    async def send_command(self, token, command):
        """发送命令到指定的客户端"""
        if token in self.clients:
            try:
                await self.clients[token].send(json.dumps({
                    'type': 'command',
                    'command': command,
                    'timestamp': datetime.now().isoformat()
                }))
                return True
            except Exception as e:
                logger.error(f"发送命令失败: {e}")
                await self.unregister(token)
        return False

    async def update_config(self, token, config):
        """更新客户端配置"""
        if token in self.clients:
            try:
                await self.clients[token].send(json.dumps({
                    'type': 'config',
                    'config': config,
                    'timestamp': datetime.now().isoformat()
                }))
                return True
            except Exception as e:
                logger.error(f"更新配置失败: {e}")
                await self.unregister(token)
        return False

    async def handle_client(self, websocket, path):
        """处理客户端连接"""
        token = None
        try:
            async for message in websocket:
                data = json.loads(message)
                msg_type = data.get('type')

                if msg_type == 'auth':
                    # 处理身份验证
                    token = data.get('token')
                    if token:
                        await self.register(websocket, token)
                
                elif msg_type == 'heartbeat':
                    # 更新心跳时间
                    if token in self.app_status:
                        self.app_status[token] = datetime.now()

        except Exception as e:
            logger.error(f"处理客户端消息失败: {e}")
        finally:
            if token:
                await self.unregister(token)

    async def start_ws_server(self, host='127.0.0.1', port=8765):
        """启动WebSocket服务器"""
        self.ws_server = await websockets.serve(self.handle_client, host, port)
        logger.info(f"WebSocket服务器启动在 ws://{host}:{port}")
        await self.ws_server.wait_closed()

    def start_websocket_server(self):
        """在新线程中启动WebSocket服务器"""
        async def run_server():
            await self.start_ws_server()

        def run_in_thread():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(run_server())
            loop.close()

        if not self.running:
            self.ws_thread = threading.Thread(target=run_in_thread, daemon=True)
            self.ws_thread.start()
            self.running = True

    def stop_websocket_server(self):
        """停止WebSocket服务器"""
        if self.running:
            if self.ws_server:
                self.ws_server.close()
            self.running = False
            if self.ws_thread:
                self.ws_thread.join(timeout=5)
                self.ws_thread = None
            logger.info("WebSocket服务器已停止")

    def start(self):
        """启动所有服务"""
        self.start_websocket_server()

    def stop(self):
        """停止所有服务"""
        self.stop_websocket_server()

# 创建全局实例
server = Server()